//  SURFObjectDetector.cpp
//
//  OpenCV Wrapper functions for Mayhem SURF-Based object detector
//
// (c) 2011 Microsoft Applied Sciences Group
// Author: Sven Kratz

#include "stdafx.h"
#include "OpenCVDLL.h"

using namespace OpenCVDLL;

/* find matching object's corners */
int SURFObjectDetector::findLastObjectCorners(int * cornerPoints)
{
	double h[9];
    CvMat _h = cvMat(3, 3, CV_64F, h);
	vector<CvPoint2D32f> pt1, pt2;
	CvMat _pt1, _pt2;
    int i, n;

	CvPoint dst_corners[4];
	CvPoint src_corners[4] = {{0,0}, {currentImage->width,0}, {currentImage->width, currentImage->height}, {0, currentImage->height}};

	// do we have enough pairs? 
	int pair_thresh = 4;
    n = (int)(ptpairs->size()/2);
    if( n < pair_thresh )
        return 0;

    pt1.resize(n);
    pt2.resize(n);
    for( i = 0; i < n; i++ )
    {
        pt1[i] = ((CvSURFPoint*) cvGetSeqElem(templateKeypoints, (*ptpairs)[i*2]))->pt;
        pt2[i] = ((CvSURFPoint*) cvGetSeqElem(imageKeypoints, (*ptpairs)[i*2+1]))->pt;
    }

    _pt1 = cvMat(1, n, CV_32FC2, &pt1[0] );
    _pt2 = cvMat(1, n, CV_32FC2, &pt2[0] );
    if( !cvFindHomography( &_pt1, &_pt2, &_h, CV_RANSAC, 3 ))
        return 0;

     for( i = 0; i < 4; i++ )
    {
        double x = src_corners[i].x, y = src_corners[i].y;
        double Z = 1./(h[6]*x + h[7]*y + h[8]);
        double X = (h[0]*x + h[1]*y + h[2])*Z;
        double Y = (h[3]*x + h[4]*y + h[5])*Z;
        dst_corners[i] = cvPoint(cvRound(X), cvRound(Y));
    }

	
	// copy to the array of ints
	int k = 0; 
	for (i = 0; i < 4; i++)
	{
		CvPoint corner = dst_corners[i];
		cornerPoints[k++] = corner.x;
		cornerPoints[k++] = corner.y;
	}

    return 1;

}

/* find pairs of matching points using fast approximate nearest neighbor (FLANN) */
void SURFObjectDetector::flannFindPairs(	const CvSeq*, 
									const CvSeq* objectDescriptors,
									const CvSeq*, 
									const CvSeq* imageDescriptors, 
									vector<int> *  ptpairs )
{
	int length = (int)(objectDescriptors->elem_size/sizeof(float));

    cv::Mat m_object(objectDescriptors->total, length, CV_32F);
	cv::Mat m_image(imageDescriptors->total, length, CV_32F);


	// copy descriptors
    CvSeqReader obj_reader;
	float * obj_ptr = m_object.ptr<float>(0);
    cvStartReadSeq( objectDescriptors, &obj_reader );
    for(int i = 0; i < objectDescriptors->total; i++ )
    {
        const float* descriptor = (const float*)obj_reader.ptr;
        CV_NEXT_SEQ_ELEM( obj_reader.seq->elem_size, obj_reader );
        memcpy(obj_ptr, descriptor, length*sizeof(float));
        obj_ptr += length;
    }
    CvSeqReader img_reader;
	float* img_ptr = m_image.ptr<float>(0);
    cvStartReadSeq( imageDescriptors, &img_reader );
    for(int i = 0; i < imageDescriptors->total; i++ )
    {
        const float* descriptor = (const float*)img_reader.ptr;
        CV_NEXT_SEQ_ELEM( img_reader.seq->elem_size, img_reader );
        memcpy(img_ptr, descriptor, length*sizeof(float));
        img_ptr += length;
    }

    // find nearest neighbors using FLANN
    cv::Mat m_indices(objectDescriptors->total, 2, CV_32S);
    cv::Mat m_dists(objectDescriptors->total, 2, CV_32F);
    cv::flann::Index flann_index(m_image, cv::flann::KDTreeIndexParams(4));  // using 4 randomized kdtrees
    flann_index.knnSearch(m_object, m_indices, m_dists, 2, cv::flann::SearchParams(64) ); // maximum number of leafs checked

    int* indices_ptr = m_indices.ptr<int>(0);
    float* dists_ptr = m_dists.ptr<float>(0);

	float thr = 0.55f;

    for (int i=0;i<m_indices.rows;++i) {
    	if (dists_ptr[2*i]<thr*dists_ptr[2*i+1]) {
    		ptpairs->push_back(i);
    		ptpairs->push_back(indices_ptr[2*i]);
    	}
    }
}

/*  object location information: calculates a homography between original image feature points and object points */
int SURFObjectDetector::locatePlanarObject( const CvSeq* objectKeypoints, 
									const CvSeq* objectDescriptors,
									const CvSeq* imageKeypoints, 
									const CvSeq* imageDescriptors,
									const CvPoint src_corners[4], CvPoint dst_corners[4] )
{
	 double h[9];
    CvMat _h = cvMat(3, 3, CV_64F, h);
    vector<int> ptpairs;
    vector<CvPoint2D32f> pt1, pt2;
    CvMat _pt1, _pt2;
    int i, n;

    flannFindPairs( objectKeypoints, objectDescriptors, imageKeypoints, imageDescriptors, &ptpairs );

	int pair_thresh = 4;
    n = (int)(ptpairs.size()/2);
    if( n < pair_thresh )
        return 0;

    pt1.resize(n);
    pt2.resize(n);

    for( i = 0; i < n; i++ )
    {
        pt1[i] = ((CvSURFPoint*)cvGetSeqElem(objectKeypoints,ptpairs[i*2]))->pt;
        pt2[i] = ((CvSURFPoint*)cvGetSeqElem(imageKeypoints,ptpairs[i*2+1]))->pt;
    }

    _pt1 = cvMat(1, n, CV_32FC2, &pt1[0] );
    _pt2 = cvMat(1, n, CV_32FC2, &pt2[0] );
    if( !cvFindHomography( &_pt1, &_pt2, &_h, CV_RANSAC, 3 ))
        return 0;

    for( i = 0; i < 4; i++ )
    {
        double x = src_corners[i].x, y = src_corners[i].y;
        double Z = 1./(h[6]*x + h[7]*y + h[8]);
        double X = (h[0]*x + h[1]*y + h[2])*Z;
        double Y = (h[3]*x + h[4]*y + h[5])*Z;
        dst_corners[i] = cvPoint(cvRound(X), cvRound(Y));
    }

    return 1;

}

/*
	calculate SURF features for the template image 
	assume 8bpp, rgb as input image

	// TODO support a grayscale format!
	// TODO report errors in some way
*/ 

void SURFObjectDetector::AddTemplate(int t_width, int t_height, unsigned char * t_imageData) 
{
	CvSize t_imageSize = cvSize(t_width, t_height); 
	IplImage * rgbImage = cvCreateImage(t_imageSize,8,3); 
	rgbImage->imageData = (char *) t_imageData; 

	// convert to grayscale image
	IplImage * tImage = cvCreateImage(t_imageSize, IPL_DEPTH_8U,1);
	cvCvtColor(rgbImage, tImage, CV_RGB2GRAY);

	// TODO: further post-processing (i.e. thresholding or histogram linearization) to improve the template image

	// calculate SURF descriptors
	
	cvClearMemStorage(templateDataStorage); 
	
	
	CvSURFParams params = cvSURFParams(SURF_HESSIAN_THRESH, 1);

	CvSeq * t_Keypoints = 0, * t_Descriptors = 0;

	cvExtractSURF( tImage, 0, &t_Keypoints, &t_Descriptors, templateDataStorage, params );

	// save references
	templateImage = tImage;
	templateKeypoints = t_Keypoints;
	templateDescriptors = t_Descriptors;

	int nTemplateKeyPoints = templateKeypoints->total; 
	int nTemplateDescriptors = templateDescriptors->total;

	//cvNamedWindow("Template",1);
	//cvShowImage("Template", tImage);


	// release memory
	cvReleaseImage(&rgbImage);
	
}

/* image processing step: calculate SURF Feature descriptors of new image */
void SURFObjectDetector::ProcessFrame(unsigned char * imageData) 
{
	// TODO
	fprintf(stdout, "SURFObjectDetector::ProcessFrame");
	currentImage->imageData  = (char *) imageData; 

	// clear memory storage from the previous iteration 
	// ----- mutex for assignment of new data to the descriptor pointers
	mut->WaitOne();
		cvClearMemStorage(storage); 

		CvSize imageSize = cvSize(img_w, img_h);

		// do grayscale conversion
		IplImage * grayImage = cvCreateImage(imageSize, IPL_DEPTH_8U,1);
		cvCvtColor(currentImage, grayImage, CV_RGB2GRAY);

		// do SIFT Feature Extraction

		CvSeq * i_Keypoints = 0, * i_Descriptors = 0;
		CvSURFParams params = cvSURFParams(SURF_HESSIAN_THRESH, 1);
	 
		cvExtractSURF(grayImage,0, &i_Keypoints, &i_Descriptors, storage, params);


		imageKeypoints = i_Keypoints;
		imageDescriptors = i_Descriptors;

		vector<int> * _ptpairs; 

		// do correspondence matching if a template image has been initialized
		if (templateImage) 
		{
			_ptpairs = new vector<int>;
		
			flannFindPairs( templateKeypoints, templateDescriptors, i_Keypoints, i_Descriptors, _ptpairs );
		}


		
		if (templateImage)
		{
			if (ptpairs)
				{
					delete ptpairs;
				}
			ptpairs = _ptpairs; 

		}
	mut->ReleaseMutex(); 
	// ----- end critical section

	cvReleaseImage(&grayImage);

#ifdef CHECK_MEMORY
	_CrtDumpMemoryLeaks();
#endif

}

/* retrieve keypoints and matching indices after image processing
   -- only works if a template image has been set
*/
void SURFObjectDetector::getKeypointsAndMatches(SURFKeyPoint * tKeyPoints, 
							int & numTemplateKeyPoints,
							SURFKeyPoint * iKeyPoints, 
							int & numImageKeyPoints, 
							int * pPairIndices, 
							int & numPointPairIndices)
{

	if (templateImage) 
	{
		// --------------- lock the buffers
		mut->WaitOne();

			// number of elements in the structs
			int nTemplateKeyPoints = templateKeypoints->total; 
			int nImageKeyPoints = imageKeypoints->total; 
			int nPointPairIndices = ptpairs->size(); 

			// write back values to reference calls
			numTemplateKeyPoints = nTemplateKeyPoints;
			numImageKeyPoints = nImageKeyPoints; 
			numPointPairIndices = nPointPairIndices;


			// copy data to buffers ------- 

			// image key points
			int idx = 0; 
			for (idx = 0; idx < nTemplateKeyPoints &&  idx < BUF_SIZE; idx++)
			{
				CvSURFPoint * pt = (CvSURFPoint*) cvGetSeqElem(templateKeypoints, idx);
				tKeyPoints[idx].x = pt->pt.x;
				tKeyPoints[idx].y = pt->pt.y;
				tKeyPoints[idx].laplacian = pt->laplacian;
				tKeyPoints[idx].size = pt->size;
				tKeyPoints[idx].dir = pt->dir;
				tKeyPoints[idx].hessian = pt->hessian;
			}

			// template keypoints
	
			for (idx = 0; idx < nImageKeyPoints &&  idx < BUF_SIZE; idx++)
			{
				CvSURFPoint * pt = (CvSURFPoint*) cvGetSeqElem(imageKeypoints, idx);
				iKeyPoints[idx].x = pt->pt.x;
				iKeyPoints[idx].y = pt->pt.y;
				iKeyPoints[idx].laplacian = pt->laplacian;
				iKeyPoints[idx].size = pt->size;
				iKeyPoints[idx].dir = pt->dir;
				iKeyPoints[idx].hessian = pt->hessian;
			}

			// point pairs -- the vector indices i and i+1 are paired image keypoints

			for (idx = 0; idx < nPointPairIndices &&  idx < 2*BUF_SIZE; idx+=2)
			{
				pPairIndices[idx] = (*ptpairs)[idx]; 
				pPairIndices[idx+1] = (*ptpairs)[idx+1]; 
			}

		// ------------------- release buffers
		mut->ReleaseMutex(); 

	}
}


SURFObjectDetector::SURFObjectDetector(int width, int height) 
{
	img_w = width;
	img_h = height; 
	CvSize imageSize =  cvSize(width, height); 
	currentImage = cvCreateImage(imageSize, 8, 3); 

	// memstorages
	storage = cvCreateMemStorage(0);
	templateDataStorage = cvCreateMemStorage(0);
}



